import gradio as gr
import cv2
import numpy as np
from tabulate import tabulate 
import torch
import os
import torch.nn as nn
from torchvision import transforms

class Perceptron(nn.Module):
    def __init__(self):
        super(Perceptron, self).__init__()
        self.flatten = nn.Flatten()  
        self.fc1 = nn.Linear(28*28, 20)  
        self.relu = nn.ReLU()  
        self.fc2 = nn.Linear(20, 10)  

    def forward(self, x):
        x = self.flatten(x)  
        x = self.fc1(x)  
        x = self.relu(x)  
        x = self.fc2(x)  
        return x

pytorch_path = 'pytorch_path.pth'  
pytorch_size = os.path.getsize(pytorch_path)  
print(f"模型参数文件路径: {pytorch_path}")  
print(f"模型参数文件大小: {pytorch_size} bytes") 

loaded_params = torch.load(pytorch_path)  
loaded_model_params = Perceptron() 
loaded_model_params.load_state_dict(loaded_params)  
loaded_model_params.eval()  

def predict(x):
    #将图像转换为张量并归一化
    preprocess = transforms.Compose([
        transforms.ToTensor(),  #将图像转变为张量
        transforms.Normalize((0.1307,), (0.3081,))  #归一化，标准化
    ])
    x_infer = preprocess(x)
    with torch.no_grad():  # 关闭梯度计算,
        y_infer = loaded_model_params(x_infer.unsqueeze(0))  
        _, predicted = torch.max(y_infer.data, 1)    
    y_max = predicted.item()
    return str(y_max)

interface = gr.Interface(predict, 
                         inputs='sketchpad', 
                         outputs='text',
                         live=True)
# 启动Gradio接口，用户可以通过这个接口进行交互
interface.launch()
